import re
from config import DESENSE_METHODS, DEFAULT_REPLACE_CHAR
import fitz  # PyMuPDF
from utils.logger import logger
from utils.ner_utils import WHITELIST

# 添加正则表达式模式用于过滤和处理
# 数字和点的任意组合模式，如1.1, 2.3.4, 1.2., .1.2等
SECTION_NUMBER_PATTERN = re.compile(
    r'''
    (?!^\s*\d{4}\.\d{1,2}(?:\.\d{1,2})?\s*$)  # 排除纯年月或年月日格式（允许前后空格）
    (?=.*[0-9])                               # 至少一个数字
    (?=.*\.)                                  # 至少一个点
    [0-9. ]+                                  # 仅数字、点、空格，且非空
    $
    ''', re.X
)

# 单个一位数字模式
SINGLE_DIGIT_PATTERN = re.compile(r'^\d$')
# 日期格式模式，匹配如"2025年7月1日"
DATE_PATTERN = re.compile(r'^\d+年\d+月\d+日$|^\d+年\d+月$|^\d+月\d+日$|^\d+年$|^\d+月$|^\d+日$')
# 提取日期中的数字
DATE_NUMBER_PATTERN = re.compile(r'\d+')

# 章节标题模式，匹配如"第一章"、"第二节"、"第十七章"等
CHAPTER_TITLE_PATTERN = re.compile(r'^第[一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾佰仟\d]+[章节篇部分编条款项目]$|^[一二三四五六七八九十百千万亿零壹贰叁肆伍陆柒捌玖拾佰仟]+[、]|^\d+[、.．]')

def filter_and_process_subgroups(char_groups):
    """
    过滤和处理子组，实现以下功能：
    1. 过滤掉"数字+'.'"的组合（如"1.""1.1""2.2.2"等）
    2. 过滤掉单个一位数字（如"1""2""6"等）
    3. 对"x年x月x日"这样的日期格式进行特殊处理，只遮挡数字部分
    4. 过滤掉"第一章"、"第二节"等章节标题
    
    Args:
        char_groups: 字符组列表
        
    Returns:
        处理后的字符组列表和需要特殊处理的日期信息
    """
    filtered_groups = []
    
    for group in char_groups:
        if not group:
            continue
            
        # 提取子组文本内容
        sub_group_text = ''.join([item.get('text', '') for item in group])
        
        # 1. 过滤掉"数字+'.'"的组合
        if SECTION_NUMBER_PATTERN.match(sub_group_text):
            logger.info(f"过滤章节编号: '{sub_group_text}'")
            continue
            
        # 2. 过滤掉单个一位数字
        if SINGLE_DIGIT_PATTERN.match(sub_group_text):
            logger.info(f"过滤单个数字: '{sub_group_text}'")
            continue
            
        # 4. 过滤掉"第一章"、"第二节"等章节标题
        if CHAPTER_TITLE_PATTERN.match(sub_group_text):
            logger.info(f"过滤章节标题: '{sub_group_text}'")
            continue
            
        # 3. 处理日期格式
        if DATE_PATTERN.match(sub_group_text):
            logger.info(f"处理日期格式: '{sub_group_text}'")
            
            # 创建数字部分子组和非数字部分子组
            digit_groups = []
            current_digit_group = []
            current_non_digit_group = []
            
            # 遍历每个字符项
            for item in group:
                item_text = item.get('text', '')
                
                # 检查字符是否为数字
                if item_text.isdigit():
                    # 如果有非数字组，先保存
                    if current_non_digit_group:
                        current_digit_group.append(item)
                    else:
                        current_digit_group.append(item)
                else:
                    # 如果有数字组，先保存
                    if current_digit_group:
                        digit_groups.append(current_digit_group)
                        current_digit_group = []
            
            # 保存最后一个数字组
            if current_digit_group:
                digit_groups.append(current_digit_group)
            
            # 将所有数字组添加到过滤后的组中
            for digit_group in digit_groups:
                if digit_group:
                    digits_text = ''.join([item.get('text', '') for item in digit_group])
                    logger.info(f"提取日期中的数字部分: '{digits_text}'")
                    filtered_groups.append(digit_group)
            
            # 跳过原始日期组
            continue
        
        # 正常添加不需要特殊处理的组
        filtered_groups.append(group)
    
    return filtered_groups

def is_strict_whitelist_match(text):
    """
    严格的白名单匹配：只有当文本完全等于白名单项或者是白名单项的子集时，才返回True
    
    Args:
        text: 要检查的文本
        
    Returns:
        bool: 是否符合严格白名单匹配条件
    """
    if not text or not text.strip():
        return False
        
    # 检查是否在主白名单中（严格匹配）
    for item in WHITELIST:
        # 只有文本完全等于白名单项或者是白名单项的子集时，才返回True
        if text == item or text in item:
            logger.info(f"严格白名单匹配: '{text}' 匹配到 '{item}'")
            return True
    
    return False

def apply_strict_whitelist_filter(char_groups):
    """
    应用严格的白名单过滤，只有当前实体内容完全等于白名单或者是白名单的子集，才算作符合白名单
    
    Args:
        char_groups: 字符组列表
        
    Returns:
        过滤后的字符组列表
    """
    filtered_groups = []
    
    for group in char_groups:
        if not group:
            continue
            
        # 提取子组文本内容
        sub_group_text = ''.join([item.get('text', '') for item in group])
        
        # 应用严格白名单匹配
        if is_strict_whitelist_match(sub_group_text):
            logger.info(f"白名单过滤: '{sub_group_text}'")
            continue
        
        # 保留不在白名单中的组
        filtered_groups.append(group)
    
    return filtered_groups

def mask_text(text, replace_char=None):
    """
    将文本替换为指定字符重复的形式
    
    Args:
        text: 要替换的原始文本
        replace_char: 替换字符，默认为配置中的DEFAULT_REPLACE_CHAR
        
    Returns:
        替换后的文本
    """
    if not replace_char:
        replace_char = DEFAULT_REPLACE_CHAR
    
    # 确保replace_char是单个字符
    if replace_char and len(replace_char) > 0:
        return replace_char[0] * len(text)
    else:
        return 'X' * len(text)  # 如果没有指定或为空，则使用X

def apply_text_desense(page, hits, method='replace', replace_char=None):
    """
    对数字 PDF 的文本块先用白框覆盖，然后再替换为指定字符
    hits: list of {'rect':(x0,y0,x1,y1),'label':..,'text':..}
    method: 脱敏方法，默认为'replace'
    replace_char: 替换字符，默认使用配置中的DEFAULT_REPLACE_CHAR
    
    Returns:
        bool: 是否有内容被处理（True表示有内容被处理，False表示所有内容都被过滤）
    """
    # 检测组内是否有换行或水平距离过大，并拆分成多个子组
    char_groups = detect_character_groups(hits)
    
    logger.info(f"\n===== 处理文本型敏感信息 =====")
    original_group_text = ''.join([item.get('text', '') for item in hits])
    
    # 收集敏感信息标签
    labels = set()
    for item in hits:
        if 'label' in item:
            labels.add(item['label'])
    
    logger.info(f"原始敏感文本: '{original_group_text}' [标签: {', '.join(labels)}]")
    logger.info(f"脱敏方法: {method}")
    
    # 过滤和处理子组
    filtered_char_groups = filter_and_process_subgroups(char_groups)
    if len(filtered_char_groups) != len(char_groups):
        logger.info(f"过滤后子组数量: {len(filtered_char_groups)}/{len(char_groups)}")
    
    # 应用严格白名单过滤
    original_filtered_count = len(filtered_char_groups)
    filtered_char_groups = apply_strict_whitelist_filter(filtered_char_groups)
    if len(filtered_char_groups) != original_filtered_count:
        logger.info(f"严格白名单过滤后子组数量: {len(filtered_char_groups)}/{original_filtered_count}")
    
    # 如果所有子组都被过滤，则不需要处理
    if len(filtered_char_groups) == 0:
        logger.info("所有内容都被过滤，不进行脱敏处理")
        return False
    
    # 处理每个子组
    for i, sub_group in enumerate(filtered_char_groups):
        if not sub_group:
            continue
            
        # 提取子组文本内容
        sub_group_text = ''.join([item.get('text', '') for item in sub_group])
        
        # 收集子组敏感信息标签
        sub_labels = set()
        for item in sub_group:
            if 'label' in item:
                sub_labels.add(item['label'])
                
        logger.info(f"处理子组 {i+1}/{len(filtered_char_groups)}: '{sub_group_text}' [标签: {', '.join(sub_labels)}]")
        
        # 为整个子组创建一个边界矩形
        x_min = float('inf')
        y_min = float('inf')
        x_max = float('-inf')
        y_max = float('-inf')
        
        # 计算包含所有字符的边界矩形
        for item in sub_group:
            rect = item['rect']
            x_min = min(x_min, rect[0])
            y_min = min(y_min, rect[1])
            x_max = max(x_max, rect[2])
            y_max = max(y_max, rect[3])
        
        # 添加内边距
        padding_x = 0.5  # 水平方向内边距，增加到3以确保左右更宽
        padding_y = 0  # 垂直方向内边距
        
        x_min -= padding_x
        y_min -= padding_y
        x_max += padding_x
        y_max += padding_y
        
        group_rect = (x_min, y_min, x_max, y_max)
        
        # 步骤1: 先用白色矩形覆盖区域
        shape = page.new_shape()
        shape.draw_rect(group_rect)  # 绘制矩形
        shape.finish(width=0, fill=(1, 1, 1))  # 白色填充，无边框
        shape.commit()
        
        # 步骤2: 对每个字符应用文本替换
        if method == 'replace':
            for item in sub_group:
                new_text = mask_text(item['text'], replace_char)
                # 使用PyMuPDF的redact功能在精确位置替换文本
                page.add_redact_annot(item['rect'], text=new_text)
    
    # 应用所有的替换
    page.apply_redactions()
    return True

def detect_character_groups(items):
    """
    检测一组字符中的分割点，将其拆分为多个合理的组
    分割标准：
    1. 垂直距离过大（换行）
    2. 水平距离过大（同一行但相距较远）
    
    Args:
        items: 字符列表（保持原始顺序）
        
    Returns:
        list: 拆分后的多个字符组列表
    """
    if not items or len(items) <= 1:
        return [items]
    
    # 计算平均字符高度和宽度，用于设置分组阈值
    heights = []
    widths = []
    for item in items:
        rect = item['rect']
        height = rect[3] - rect[1]  # y1 - y0
        width = rect[2] - rect[0]   # x1 - x0
        heights.append(height)
        widths.append(width)
    
    avg_height = sum(heights) / len(heights)
    avg_width = sum(widths) / len(widths)
    
    # 分别设置垂直和水平方向的阈值
    vertical_threshold = avg_height * 0.6  # 垂直方向阈值
    horizontal_threshold = avg_width * 4.0  # 水平方向阈值
    
    # 拆分为多组
    groups = []
    current_group = [items[0]]
    
    for i in range(1, len(items)):
        prev = items[i-1]
        curr = items[i]
        
        # 获取当前和前一个字符的位置信息
        prev_rect = prev['rect']
        curr_rect = curr['rect']
        
        # 计算水平和垂直方向的距离
        # 使用右边界到左边界的距离作为水平间距
        horizontal_distance = curr_rect[0] - prev_rect[2]
        
        # 使用中心点的垂直距离作为垂直间距
        prev_center_y = (prev_rect[1] + prev_rect[3]) / 2
        curr_center_y = (curr_rect[1] + curr_rect[3]) / 2
        vertical_distance = abs(curr_center_y - prev_center_y)
        
        # 如果水平或垂直距离超过对应阈值，则认为需要分组
        if horizontal_distance > horizontal_threshold or vertical_distance > vertical_threshold:
            # 需要新建一个组
            if current_group:
                groups.append(current_group)
            current_group = [curr]
        else:
            # 添加到当前组
            current_group.append(curr)
    
    # 添加最后一个组
    if current_group:
        groups.append(current_group)
    
    return groups

def apply_image_desense(page, hits, method='mask', mask_color=None, mosaic_size=None, replace_char=None):
    """
    对图像型 PDF 在坐标上进行多种脱敏处理
    method: 'mask'|'blur'|'replace'|'inpaint'
    mask_color: 遮罩颜色的十六进制值，如 "#000000"（黑色）
    mosaic_size: 马赛克块大小，默认为10像素
    replace_char: 替换字符，默认使用配置中的DEFAULT_REPLACE_CHAR
    
    对连续字符进行整体遮挡，而不是单个字符遮挡：
    1. 高度统一为这组字符中位数高度
    2. 检测换行并拆分为多组处理
    3. 第一个字符向右侧加宽，最后一个字符向左侧加宽，中间字符向两侧加宽
    
    Returns:
        bool: 是否有内容被处理（True表示有内容被处理，False表示所有内容都被过滤）
    """
    from config import MASK_COLOR, DEFAULT_MOSAIC_SIZE, DEFAULT_REPLACE_CHAR
    
    # 如果未指定mask_color，则使用配置中的默认值
    if not mask_color and method == 'mask':
        mask_color = MASK_COLOR
        
    # 如果未指定mosaic_size，则使用配置中的默认值
    if not mosaic_size and method == 'blur':
        mosaic_size = DEFAULT_MOSAIC_SIZE
    else:
        # 确保mosaic_size是整数
        try:
            mosaic_size = int(mosaic_size)
        except (TypeError, ValueError):
            mosaic_size = DEFAULT_MOSAIC_SIZE
            
    # 如果未指定replace_char，则使用配置中的默认值
    if not replace_char and method == 'replace':
        replace_char = DEFAULT_REPLACE_CHAR

    # 将十六进制颜色转换为RGB（0-1范围）
    mask_rgb = (0, 0, 0)  # 默认黑色
    if mask_color and mask_color.startswith('#') and len(mask_color) == 7:
        try:
            r = int(mask_color[1:3], 16) / 255.0
            g = int(mask_color[3:5], 16) / 255.0
            b = int(mask_color[5:7], 16) / 255.0
            mask_rgb = (r, g, b)
        except ValueError:
            logger.info(f"无效的颜色值 {mask_color}，使用默认黑色")
            
    # 根据group_id对hits分组
    groups = {}
    for item in hits:
        group_id = item.get('group_id', None)
        if group_id:
            if group_id not in groups:
                groups[group_id] = []
            groups[group_id].append(item)
    
    logger.info(f"\n===== 处理图像型敏感信息 =====")
    logger.info(f"脱敏方法: {method}")
    logger.info(f"敏感信息组数: {len(groups)}")
    
    # 对每个分组处理，创建新的矩形区域
    merged_rects = []
    
    for group_id, items in groups.items():
        if not items:
            continue
            
        # 不再按X坐标排序，保留原始检测顺序
        # 检测组内是否有换行或水平距离过大，并拆分成多个子组
        char_groups = detect_character_groups(items)
        
        # 记录组内容到日志
        original_group_text = ''.join([item.get('text', '') for item in items])
        
        # 收集敏感信息标签
        labels = set()
        for item in items:
            if 'label' in item:
                labels.add(item['label'])
        
        logger.info(f"\n处理敏感信息组: '{original_group_text}' [标签: {', '.join(labels)}]")
        
        # 过滤和处理子组
        filtered_char_groups = filter_and_process_subgroups(char_groups)
        if len(filtered_char_groups) != len(char_groups):
            logger.info(f"过滤后子组数量: {len(filtered_char_groups)}/{len(char_groups)}")
        
        # 应用严格白名单过滤
        original_filtered_count = len(filtered_char_groups)
        filtered_char_groups = apply_strict_whitelist_filter(filtered_char_groups)
        if len(filtered_char_groups) != original_filtered_count:
            logger.info(f"严格白名单过滤后子组数量: {len(filtered_char_groups)}/{original_filtered_count}")
        
        # 如果该组的所有子组都被过滤，则跳过该组
        if len(filtered_char_groups) == 0:
            logger.info(f"组 '{original_group_text}' 的所有内容都被过滤，跳过处理")
            continue
        
        # 处理每个子组
        for i, sub_group in enumerate(filtered_char_groups):
            # 提取子组文本内容
            sub_group_text = ''.join([item.get('text', '') for item in sub_group])
            
            # 收集子组敏感信息标签
            sub_labels = set()
            for item in sub_group:
                if 'label' in item:
                    sub_labels.add(item['label'])
                    
            logger.info(f"处理子组 {i+1}/{len(filtered_char_groups)}: '{sub_group_text}' [标签: {', '.join(sub_labels)}]")
            
            if not sub_group:
                continue
                
            # 为整个子组创建一个边界矩形
            if len(sub_group) == 0:
                continue
                
            # 初始化边界值
            x_min = float('inf')
            y_min = float('inf')
            x_max = float('-inf')
            y_max = float('-inf')
            
            # 计算包含所有字符的边界矩形
            for item in sub_group:
                rect = item['rect']
                x_min = min(x_min, rect[0])
                y_min = min(y_min, rect[1])
                x_max = max(x_max, rect[2])
                y_max = max(y_max, rect[3])
            
            # 添加内边距
            padding_x = 0.5 # 水平方向内边距，增加到3以确保左右更宽
            padding_y = 0  # 垂直方向内边距
            
            x_min -= padding_x
            y_min -= padding_y
            x_max += padding_x
            y_max += padding_y
            
            # 记录合并后的矩形
            merged_rect = (x_min, y_min, x_max, y_max)
            #logger.info(f"  合并矩形: ({x_min:.1f},{y_min:.1f},{x_max:.1f},{y_max:.1f})")
            
            if method == 'blur':
                # 对矩形区域进行高斯模糊
                merged_rects.append((merged_rect, 'blur', None))
            elif method == 'replace':
                # 对矩形区域进行文本替换
                merged_rects.append((merged_rect, 'replace', mask_text(sub_group_text, replace_char)))
            else:  # 默认使用mask方法
                # 对矩形区域进行遮罩
                merged_rects.append((merged_rect, 'mask', mask_rgb))
    
    # 如果所有组都被过滤，没有需要处理的矩形
    if not merged_rects:
        logger.info("所有内容都被过滤，不进行脱敏处理")
        return False
        
    # 应用脱敏效果
    for rect, effect_type, effect_value in merged_rects:
        shape = page.new_shape()
        shape.draw_rect(rect)  # 绘制矩形
        
        if effect_type == 'mask':
            # 使用自定义颜色填充矩形
            shape.finish(width=0, fill=effect_value)
        elif effect_type == 'blur':
            # 实现马赛克效果
            # 获取矩形区域的坐标
            x0, y0, x1, y1 = rect
            
            # 截取矩形区域的图像
            pix = page.get_pixmap(clip=fitz.Rect(rect))
            
            # 将pixmap转换为numpy数组/OpenCV图像
            import numpy as np
            import cv2
            
            img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, pix.n)
            img = np.copy(img)  # 创建可写副本
            
            # 应用马赛克效果
            # mosaic_size = 10  # 可调整马赛克块大小
            
            # 获取图像尺寸
            h, w = img.shape[:2]
            
            # 将图像划分为马赛克块
            for y in range(0, h, mosaic_size):
                for x in range(0, w, mosaic_size):
                    # 确保不超出图像边界
                    y_end = min(y + mosaic_size, h)
                    x_end = min(x + mosaic_size, w)
                    
                    # 获取当前马赛克块
                    block = img[y:y_end, x:x_end]
                    
                    # 计算整个块的平均颜色值
                    if block.size > 0:  # 确保块非空
                        # 对于彩色图像
                        if len(block.shape) == 3:
                            avg_color = block.mean(axis=(0, 1))
                        # 对于灰度图像
                        else:
                            avg_color = block.mean()
                        
                        # 用平均颜色填充整个块
                        block[:] = avg_color
            
            # 创建临时文件保存马赛克处理后的图像
            import tempfile, os
            temp_fd, temp_path = tempfile.mkstemp(suffix='.png')
            os.close(temp_fd)
            
            # 保存马赛克后的图像
            if pix.n == 3 or pix.n == 4:  # RGB或RGBA
                img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            else:  # 灰度图
                img_bgr = img
            
            cv2.imwrite(temp_path, img_bgr)
            
            # 将马赛克后的图像插入回PDF
            page.insert_image(fitz.Rect(rect), filename=temp_path)
            
            # 删除临时文件
            os.remove(temp_path)
            
            continue  # 跳过下面的shape.commit()
        elif effect_type == 'replace':
            # 文本替换效果
            # 首先清除区域
            shape.finish(width=0, fill=(1, 1, 1))
            shape.commit()
            
            # 然后添加替换文本（实际上是文字注解）
            page.insert_text((rect[0], rect[1]), effect_value or "XXXX", fontsize=12, color=(0, 0, 0))
            continue  # 跳过下面的commit
            
        shape.commit()
    
    return True  # 返回True表示成功处理了内容